import matplotlib.pyplot as plt
import numpy as np
from matplotlib.ticker import MaxNLocator, MultipleLocator
import sys
from os.path import abspath, dirname

plt.rc('font', family='Arial', weight='normal')
if len(sys.argv) <= 1:
    print("using parm like dcqcn or hp95")
    exit(-1)
cc = sys.argv[1]

dir_in = 'output/'
dir_out = 'img/'
if "backup" in dirname(abspath(__file__)):
    dir_in = ''
    dir_out = ''

# 初始化流数据字典
flows = {}

# 读取流完成时间
fcts = []
with open(f'{dir_in}fct_{cc}_test.txt', 'r') as f:
    for line in f:
        parts = line.split()
        start_time = int(parts[5]) - 3000000000
        fct = int(parts[6])
        
        fcts.append(start_time + fct)
        
        
# 读取文件并筛选数据
with open(f'{dir_in}send_rate_{cc}_test.txt', 'r') as f:
    for line in f:
        parts = line.split()
        src_node = int(parts[0])
        dst_node = int(parts[1])
        src_port = int(parts[2])
        timestamp = int(parts[5]) - 3000000000
        throughput = int(parts[4])

        if throughput == 0:
            continue
        
        # 使用 (src_node, dst_node, src_port) 作为流的唯一标识
        flow_key = (src_node, dst_node, src_port)
        
        if flow_key not in flows:
            # 如果流还没有在字典中，初始化该流的数据
            flows[flow_key] = {'x': [], 'y': []}
        
        flows[flow_key]['x'].append(timestamp)
        flows[flow_key]['y'].append(throughput)

# 降采样：每隔10个数据点选择一个数据点
sample_interval = 1
for flow_key in flows:
    flows[flow_key]['x'] = flows[flow_key]['x'][::sample_interval] + [flows[flow_key]['x'][-1]]
    flows[flow_key]['y'] = flows[flow_key]['y'][::sample_interval] + [flows[flow_key]['y'][-1]]
    
    # 转换单位，时间以秒为单位，吞吐量以Gbps为单位
    flows[flow_key]['x'] = np.array(flows[flow_key]['x']) / 1e6  # 转换为秒
    flows[flow_key]['y'] = np.array(flows[flow_key]['y']) / 1e9  # 转换为Gbps

    # 重新采样：使X轴取值均匀
    x_resampled = np.linspace(flows[flow_key]['x'][0], flows[flow_key]['x'][-1], len(flows[flow_key]['x']))
    
    # 线性插值
    y_resampled = np.interp(x_resampled, flows[flow_key]['x'], flows[flow_key]['y'])

    # 更新流的x和y数据
    flows[flow_key]['x'] = x_resampled
    flows[flow_key]['y'] = y_resampled

# 设置绘图的样式
plt.figure(figsize=(30, 8))
plt.tick_params(labelsize=40)
font1 = {'weight': 'normal', 'size': 40}
X_name = "Time(ms)"
plt.xlabel(X_name, fontdict=font1)
Y_name = "Throughput(Gbps)"
plt.ylabel(Y_name, fontdict=font1)

# 绘制每条流的图形
colors = ['#5b7edd', '#6ba464', '#c96549', '#305089', '#a43337', '#87CEEB', '#32CD32', '#FFD700']  # 添加更多的颜色
linestyles = ['-', '-', '-']
idx = 0
for flow_key, data in flows.items():
    x_resampled = data['x']
    y_resampled = data['y']
    
    # 获取流的标签
    label = f"F{idx + 1}"
    
    # 绘制曲线
    plt.plot(x_resampled, y_resampled, label=label, linewidth=7, color=colors[idx % len(colors)], linestyle=linestyles[idx % len(linestyles)])
    idx += 1

# 添加辅助说明曲线
# for fct in fcts:
#     plt.axvline(x=fct/1e6, color='k', linestyle='-.', linewidth=5, alpha=0.7)

# 调整图像
plt.xlim([0, 42])
# plt.ylim([2, 10.5])


# 显示更多刻度
plt.xticks(np.arange(0, 42, 10))  # 设置横坐标显示的刻度
plt.yticks(np.arange(0.1, 2.1, 0.1))  # 设置纵坐标显示的刻度

# 设置纵坐标不显示小数，强制显示整数
plt.gca().yaxis.set_major_locator(MaxNLocator(integer=True))

# 启用小刻度，并显示
plt.minorticks_on()

# 自定义小刻度间隔，减少小刻度的密度
plt.gca().xaxis.set_minor_locator(MultipleLocator(5))  # 每5显示一个小刻度
plt.gca().yaxis.set_minor_locator(MultipleLocator(1))  # 每1显示一个小刻度

# 设置主刻度和小刻度的样式
plt.tick_params(which='major', axis='x', direction='in', length=10, width=4)  # 主刻度和小刻度的样式
plt.tick_params(which='major', axis='y', direction='in', length=10, width=4)
plt.tick_params(which='minor', axis='x', direction='in', length=4, width=4)  # 主刻度和小刻度的样式
plt.tick_params(which='minor', axis='y', direction='in', length=4, width=4)

# 添加图例
plt.legend(loc='lower right', ncol=1, prop={'size': 30, 'weight': 'normal'})

# 保存图像
png_filename = f'{dir_out}fig1.pdf'
plt.savefig(png_filename, format='pdf', dpi=300, bbox_inches='tight')
png_filename = f'{dir_out}fig1.png'
plt.savefig(png_filename, format='png', dpi=300, bbox_inches='tight')

print('finish plot.')